#include "cfsrbt.h"
#include "proto.h"
#include "global.h"


extern int strcmp(const char *s1, const char *s2);
// 比较函数，用于比较两个 cfs_node 结构体的 keystring
int mytype_cmp(const char *a, const char *b) {
    return strcmp(a, b);
}

// cfstree插入比较
runtime_t cfs_cmp(cfs_node* n1, cfs_node* n2) {
    return n1->vruntime - n2->vruntime;
}

int get_weight(int pri){
    return sched_prio_to_weight[pri + 20];
}

// 增加 vruntime
runtime_t add_vruntime(cfs_node* node, int ticks) {
    if (!node) {
        printf("cfs node is null\n");
        return -1;
    }

    // if (xchg(&node->lock, 1) == 1) {
    //     return -2;
    // }
    int pri = node->nice;
    runtime_t cfs_vtime = (runtime_t)((double)ticks * (double)NICE_0_WEIGHT / (double)get_weight(pri));
    // assert(cfs_vtime >= 0);
    node->vruntime += cfs_vtime;
    // node->vruntime += (int)(ticks * NICE_0_WEIGHT / get_weight(pri)) ;

    xchg(&node->lock, 0);
    return 0;
}

// 初始化cfs_node
struct cfs_node *cfs_node_init(cfs_node* node, PROCESS* proc, int nice) {
    rb_init_node(&node->node);
    node->keystring = proc->task.p_name;
    node->proc = proc;
    node->lock = 0;
    node->nice = nice;
    node->vruntime = 0;
    node->rruntime = 0;
    node->start = 0;
    return node;
}

// 创建一个CFS节点
struct cfs_node *create_cfs_node(PROCESS *proc, int nice) {
    struct cfs_node *node = (cfs_node *)malloc((u32)sizeof(struct cfs_node));
    // struct cfs_node *node = NULL;
    if (!node) {
        return NULL; // 内存分配失败
    }

    node->proc = proc;
    node->keystring = proc->task.p_name; // 假设进程名为键值
    node->lock = 0;
    node->nice = nice;
    node->vruntime = 0;
    node->rruntime = 0;
    node->start = 0;

    return node;
}

// 通过vruntime进行搜索
struct cfs_node *cfs_search(struct rb_root *root, struct cfs_node* tar) {
    struct rb_node *node = root->rb_node;

    while (node) {
        struct cfs_node *data = rb_entry(node, struct cfs_node, node);
        runtime_t result = cfs_cmp(data, tar);

        if (result > 0)
            node = node->rb_left;
        else if (result < 0)
            node = node->rb_right;
        else
            return data;
    }
    return NULL;
}


// 通过pid进行搜索操作
struct cfs_node *cfs_search_pid(struct rb_root *root, u32 pid) {
    struct rb_node *node;

    for (node = rb_first(root); node; node = rb_next(node)) {
        struct cfs_node *data = rb_entry(node, struct cfs_node, node);
        if (data->proc->task.pid == pid) {
            return data;
        }
    }
    return NULL;
}

// 插入操作
int cfs_insert(struct rb_root *root, struct cfs_node *data) {
    while (xchg(&cfs_lock, 1) == 1) {
        // 休眠
        // printf("lock\n");
        int myticks = ticks;
        while (ticks - myticks < DEFAULT_TICKS) ;
    }

    struct rb_node **new = &(root->rb_node), *parent = NULL;
    // malloc(sizeof(cfs_node));
    while (*new) {
        struct cfs_node *this = rb_entry(*new, struct cfs_node, node);
        runtime_t result = cfs_cmp(data, this);

        parent = *new;
        if (result < 0)
            new = &((*new)->rb_left);
        else if (result >= 0)
            new = &((*new)->rb_right);
        // else {
        //     xchg(&cfs_lock, 0);
        //     return FALSE; // 已存在相同字符串，插入失败
        // }
    }

    // 新节点插入到红黑树中
    rb_link_node(&data->node, parent, new);
    rb_insert_color(&data->node, root);

    xchg(&cfs_lock, 0);

    return TRUE; // 插入成功
}


// 更新节点的 vruntime 并调整红黑树结构
int cfs_update_vruntime(struct rb_root *root, struct cfs_node *node, runtime_t new_vruntime) {
    if (!node)
        return -1; // 节点不存在

    while (xchg(&cfs_lock, 1) == 1) {
        // 休眠
        int myticks = ticks;
        while (ticks - myticks < DEFAULT_TICKS) ;
    }

    while (xchg(&node->lock, 1) == 1) {
        // 休眠
        int myticks = ticks;
        while (ticks - myticks < DEFAULT_TICKS) ;
    }
    // cfstree中找不到此节点
    if (!cfs_search_pid(root, node->proc->task.pid)) 
        return -2;

    // 从树中删除节点
    rb_erase(&node->node, root);

    // 更新节点的 vruntime
    node->vruntime = new_vruntime;

    // 重新将节点插入到树中
    struct rb_node **new = &(root->rb_node), *parent = NULL;

    while (*new) {
        struct cfs_node *this = rb_entry(*new, struct cfs_node, node);
        runtime_t result = cfs_cmp(node, this);

        parent = *new;
        if (result < 0)
            new = &((*new)->rb_left);
        else if (result >= 0)
            new = &((*new)->rb_right);
        // else {
        //     xchg(&cfs_lock, 0);
        //     xchg(&node->lock, 0);
        //     return -3; // 已存在相同 vruntme 的节点
        // }
    }

    // 新节点插入到红黑树中
    rb_link_node(&node->node, parent, new);
    rb_insert_color(&node->node, root);

    xchg(&cfs_lock, 0);
    xchg(&node->lock, 0);

    return 0; // 更新成功
}



// 删除操作
int cfs_remove(struct rb_root *root, struct cfs_node *data) {
    if (!data)
        return FALSE; // 节点不存在

    while (xchg(&cfs_lock, 1) == 1) {
        // 休眠
        int myticks = ticks;
        while (ticks - myticks < DEFAULT_TICKS) ;
    }

    while (xchg(&data->lock, 1) == 1) {
        // 休眠
        int myticks = ticks;
        while (ticks - myticks < DEFAULT_TICKS) ;
    }

    struct rb_node *node = &(data->node);
    if (!node) {
        xchg(&cfs_lock, 0);
        xchg(&data->lock, 0);
        return FALSE; // 节点不存在于红黑树中
    }

    rb_erase(node, root); // 从红黑树中删除节点

    // 释放节点的内存
    // free(data->keystring); // 如果keystring是动态分配的，需要释放
    xchg(&cfs_lock, 0);
    xchg(&data->lock, 0);


    return TRUE; // 删除成功
}

// 从CFS调度器中删除进程
int cfs_remove_process(struct rb_root *root, PROCESS *proc) {
    struct cfs_node *node = cfs_search_pid(root, proc->task.pid);
    if (!node) {
        return FALSE; // 未找到节点
    }

    return cfs_remove(root, node); // 调用已有的删除函数删除节点
}


// CFS调度函数：选择下一个要执行的进程
PROCESS *cfs_schedule_next(struct rb_root *root) {
    if (!root->rb_node) return NULL;
    struct cfs_node *next_node = rb_entry(rb_first(root), struct cfs_node, node);
    if (!next_node) {
        return NULL; // 没有可执行的进程
    }

    return next_node->proc; // 返回下一个要执行的进程
}

// 更新进程的运行时间（vruntime）
runtime_t update_vruntime(struct cfs_node *node, int ticks) {
    if (!node) {
        return -1; // 节点不存在
    }

    return add_vruntime(node, ticks); // 调用已有的函数更新 vruntime
}



// 打印进程信息
void print_process_info(PROCESS *proc) {
    uart_kprintf("Process ID: %u\n", proc->task.pid);
    uart_kprintf("Process Name: %s\n", proc->task.p_name);
    // uart_kprintf("Priority: %d\n", proc->task.priority);
    // uart_kprintf("Ticks Remaining: %d\n", proc->task.ticks);
    // 打印其他字段信息...
    // uart_kprintf("\n");
}

// 遍历CFS调度器并打印进程信息
void print_cfs_scheduler(struct rb_root *root) {
    struct rb_node *node;
    uart_kprintf("CFS Scheduler Contents:\n");

    for (node = rb_first(root); node; node = rb_next(node)) {
        struct cfs_node *data = rb_entry(node, struct cfs_node, node);
        print_process_info(data->proc);
        // uart_kprintf("Process vruntime: %d\n", data->vruntime);
        uart_kprintf("Process vruntime: %d\n", (int)(data->vruntime));
        uart_kprintf("Process rruntime: %d\n\n", data->rruntime);
    }
}


// 设置nice值
void set_nice(struct cfs_node* node, int nice) {
    while (xchg(&node->lock, 1) == 1) {
        // 休眠
        int myticks = ticks;
        while (ticks - myticks < DEFAULT_TICKS) ;
    }

    if(nice < -20 || nice >19) {
        xchg(&node->lock, 0);
        return;
    }

    node->nice = nice;

    xchg(&node->lock, 0);
}

//设置vruntime值
void set_vruntime(struct cfs_node* node, runtime_t vruntime) {
    while (xchg(&node->lock, 1) == 1) {
        // 休眠
        int myticks = ticks;
        while (ticks - myticks < DEFAULT_TICKS) ;
    }

    node->vruntime = vruntime;

    xchg(&node->lock, 0);
}

// 返回cfstree中最小的vruntime
runtime_t get_min_vruntime(struct rb_root *root) {
    if (!root || !root->rb_node) return -1;
    struct rb_node *node = node = rb_first(root);
    struct cfs_node *data = rb_entry(node, struct cfs_node, node);
    return data->vruntime;
}

